///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file corr_prod.h
///
/// \brief   header file for the correlator product class
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef PCPS_CORR_PROD_HEADER
#define PCPS_CORR_PROD_HEADER

#include<src/pcps.h>
#include<src/input.h>
#include<src/correlator.h>

namespace pcps {

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \brief   Class to represent a correlator product.
  ///
  ///////////////////////////////////////////////////////////////////////////////////////////////////
  template<class S> class CorrProd : public CoeffObj<S> {

    private:

      /// \brief  For each lattice site, a vector of the coefficients of that site's occupation number
      ///         in the compound index formulas for the correlators that touch that site.
      std::vector<std::vector<int> > _occ_coeffs;

      /// \brief  For each lattice site, a vector giving the indeces of correlators touching the site.
      std::vector<std::vector<int> > _touching_corrs;

      /// \brief  For each correlator, the compound index for the most recently initialized lattice configuration.
      std::vector<int> _cmpd_inds;

      /// \brief  For each correlator, the pointer to that correlator's data
      std::vector<S *> _data_ptrs;

      /// \brief  For each correlator, the position of that correlator's data in the overall ordering of the data
      std::vector<int> _data_pos;

      /// \brief  The number of correlators with fixed values.
      int _n_fixed;

      /// \brief  A number that the magnitude of doubles should not exceed during some calculations
      const double _large;

      /// \brief  A number that the magnitude of doubles should not be smaller than during some calculations
      const double _small;

      /// \brief  Returns whether a real number has an extremely large or small magnitude
      double is_extreme(const double x) {
        return ( x > _large || x < -_large || ( x > -_small && x < _small ) );
      }

      /// \brief  Returns whether a complex number has an extremely large or small magnitude
      double is_extreme(const std::complex<double> & x) {
        return ( ( x.real() > _large || x.real() < -_large || x.imag() > _large || x.imag() < -_large ) ||
                 ( x.real() > -_small && x.real() < _small && x.imag() > -_small && x.imag() < _small )    );
      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief Updates the internal data for a one-site occupation change and multiplies the supplied
      ///        value by the ratio of the new coefficient to the old coefficient.
      ///
      /// \param[in]      userinp  user input options
      /// \param[in,out]  text     a text stream from which to construct the object.
      /// \param[in]      co_map   map of names to available coefficient objects
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void update_one_site(const int site, const int new_occ, const int old_occ, pcps::MantExp<S> & val) {

        // get the number of correlators touching the site
        const int n = _touching_corrs[site].size();

        // if no correlators touch the site, we are finished
        if (n == 0) return;

        // make sure the number of touching correlators matches the number of occupation coefficients
        assert( _touching_corrs[site].size() == _occ_coeffs[site].size() );

        // get pointers for the correlator indices and occupation coefficients
        const int * const tc = &_touching_corrs[site][0];
        const int * const occ_coeff = &_occ_coeffs[site][0];

        // compute the difference between the new and old occupations
        const int occ_diff = new_occ - old_occ;

        // initialize the sum of differences between the new and old correlator values
        S diff_sum = pcps::zero<S>();

        // declare an integer used to loop over the correlators touching the site
        int i;

        // declare an integer for a given correlator's position in the overall correlator ordering
        int c;

        // declare a variable to hold the old correlator value
        S x;

        // to encourage pipelining, update correlators 3 at a time
        for (i = 0; i < n-2; i+=3) {

          c = tc[i+0];                                   // get index of the correlator
          x = _data_ptrs[c][_cmpd_inds[c]];              // get correlator's old value
          _cmpd_inds[c] += occ_coeff[i+0] * occ_diff;    // update the correlator's compound index
          diff_sum += _data_ptrs[c][_cmpd_inds[c]] - x;  // add the new/old value difference to the total

          c = tc[i+1];                                   // get index of the correlator
          x = _data_ptrs[c][_cmpd_inds[c]];              // get correlator's old value
          _cmpd_inds[c] += occ_coeff[i+1] * occ_diff;    // update the correlator's compound index
          diff_sum += _data_ptrs[c][_cmpd_inds[c]] - x;  // add the new/old value difference to the total

          c = tc[i+2];                                   // get index of the correlator
          x = _data_ptrs[c][_cmpd_inds[c]];              // get correlator's old value
          _cmpd_inds[c] += occ_coeff[i+2] * occ_diff;    // update the correlator's compound index
          diff_sum += _data_ptrs[c][_cmpd_inds[c]] - x;  // add the new/old value difference to the total

          //// if the ratio has become very large or very small, record and reset it to prevent
          //// a floating point overflow
          //if ( this->is_extreme(ratio) ) {
          //  val = val * ratio;
          //  ratio = pcps::unity<S>();
          //}

          //// check for NaN
          //if ( ratio != ratio )
          //  throw pcps::Exception("encountered NaN (%.2e) in CorrProd::update_one_site") % std::abs(ratio);

          //// check for large ratio
          //if ( this->is_extreme(ratio) )
          //  throw pcps::Exception("encountered extreme ratio (%.2e) in CorrProd::update_one_site") % std::abs(ratio);

        }

        // update remaining correlators
        for ( ; i < n; i++) {
          c = tc[i+0];                                   // get index of the correlator
          x = _data_ptrs[c][_cmpd_inds[c]];              // get correlator's old value
          _cmpd_inds[c] += occ_coeff[i+0] * occ_diff;    // update the correlator's compound index
          diff_sum += _data_ptrs[c][_cmpd_inds[c]] - x;  // add the new/old value difference to the total
        }

        // record the coefficient ratio
        val = val * std::exp(diff_sum);

      }

    public:

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief Constructs the correlator product.
      ///
      /// \param[in]      userinp  the user's input options
      /// \param[in,out]  arch     an archive from which to read the correlator
      /// \param[in]      wi_map   map of names to available integer wavefunction data objects
      /// \param[in]      wf_map   map of names to available floating point wavefunction data objects
      /// \param[in]      co_map   map of names to available coefficient objects
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      CorrProd(const pcps::Input & userinp,
               pcps::Archive & arch,
               const std::map<std::string, boost::shared_ptr<pcps::WfnData<int> > > & wi_map,
               const std::map<std::string, boost::shared_ptr<pcps::WfnData<S> > > & wf_map,
               const std::map<std::string, boost::shared_ptr<pcps::CoeffObj<S> > > & co_map)
        : CoeffObj<S>(0, 0, 0),
          _occ_coeffs( userinp.nsites() ),
          _touching_corrs( userinp.nsites() ),
          _cmpd_inds(),
          _data_ptrs(),
          _data_pos(),
          _large( std::pow(double(10.0), std::numeric_limits<double>::max_exponent10 / 3) ),
          _small( std::pow(double(10.0), std::numeric_limits<double>::min_exponent10 / 3) )
      {

        // get name of the correlator product
        pcps::mpi::read_and_bcast(arch, this->_name, "Failed to read name in CorrProd constructor.");

        // get correlator names
        std::vector<std::string> corr_names;
        pcps::mpi::read_and_bcast(arch, corr_names, "Failed to read vector of correlator names in CorrProd constructor.");

        // add correlators to the vector of coefficient object dependencies
        for (std::vector<std::string>::const_iterator cname = corr_names.begin(); cname != corr_names.end(); cname++) {

          // check that name exists
          if ( pcps::mpi::rank() == 0 && co_map.count(*cname) == 0 )
            throw pcps::Exception("no coefficient object by the name %s") % *cname;

          // add the object to the dependency list
          this->_co_deps->push_back(co_map.find(*cname)->second);

          // check that the object is indeed a correlator
          if ( pcps::mpi::rank() == 0 && typeid(**this->_co_deps->rbegin()) != typeid(pcps::Correlator<S>) )
            throw pcps::Exception("coefficient object with name \"%s\" is not a Correlator") % *cname;

        }

        // if an empty vector was supplied, add all available correlators to the correlator product
        if ( corr_names.size() == 0 ) {
          typename std::map<std::string, boost::shared_ptr<pcps::CoeffObj<S> > >::const_iterator it;
          for (it = co_map.begin(); it != co_map.end(); it++)
            if ( typeid(*it->second) == typeid(pcps::Correlator<S>) )
              this->_co_deps->push_back(it->second);
        }

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief converts the object into the command that will create the object
      ///
      /// \param[in,out]  arch     an archive in which to store the command
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void to_command(pcps::Archive & arch) const {

        // create a vector of the names of the correlators
        std::vector<std::string> v;
        typename pcps::PtrTypes<pcps::CoeffObj<S> >::vp::const_iterator it;
        for (it = this->_co_deps->begin(); it != this->_co_deps->end(); it++)
          v.push_back((**it).name());

        // write the command
        arch << pcps::CorrProdCommand::keyword() << std::endl;
        arch << this->_name                      << std::endl;
        arch << v                                << std::endl;
        arch << std::endl;

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief returns the value of the coefficient for an arbitrary configuration
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      pcps::MantExp<S> get_value(const pcps::Configuration<S> & config) {
        throw pcps::Exception("NOT YET IMPLEMENTED");
        return pcps::MantExp<S>();
      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief returns the value of the coefficient for a configuration that differs from the
      ///        currently initialized configuration as specified.
      ///
      /// \param[in]      new_config  the lattice configuration whose coefficient we want
      /// \param[in]      old_config  the lattice configuration the object is currently initialized for
      /// \param[in]      diff        an object specifying how the configurations differ
      ///
      /// \return the coefficient at the new configuration
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      pcps::MantExp<S> get_value(const pcps::Configuration<S> & new_config,
                                 const pcps::Configuration<S> & old_config,
                                 const pcps::ConfigDiff<S> & diff) {

        for (std::vector<int>::const_iterator site = diff.newly_empty().begin(); site != diff.newly_empty().end(); site++)
          this->update_one_site(*site, 0, 1, this->_curr_val);

        for (std::vector<int>::const_iterator site = diff.newly_full().begin(); site != diff.newly_full().end(); site++)
          this->update_one_site(*site, 1, 0, this->_curr_val);

        pcps::MantExp<S> retval(this->_curr_val);

        for (std::vector<int>::const_iterator site = diff.newly_empty().begin(); site != diff.newly_empty().end(); site++)
          this->update_one_site(*site, 1, 0, this->_curr_val);

        for (std::vector<int>::const_iterator site = diff.newly_full().begin(); site != diff.newly_full().end(); site++)
          this->update_one_site(*site, 0, 1, this->_curr_val);

        return retval;

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief Computes the ratios of the noiseless coefficients at each singly excited configuration
      ///        to the noisy coefficient at the current configuration
      ///
      /// \param[in]      config      the lattice configuration the object is currently initialized for
      /// \param[out]     a_ratios    on exit, the matrix of excitation ratios for alpha excitations
      /// \param[out]     b_ratios    on exit, the matrix of excitation ratios for beta excitations
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void single_excitation_ratios(const pcps::Configuration<S> & config,
                                    std::vector<S> & a_ratios,
                                    std::vector<S> & b_ratios)
      {
        throw pcps::Exception("NOT YET IMPLEMENTED");
      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief Adds the derivatives of the coefficient with respect to the wavefunction's data
      ///        variables to the supplied vector after multiplying the derivatives by a constant.
      ///        The derivatives are computed assuming that the object is currently initialized for the
      ///        supplied lattice configuration.
      ///
      /// \param[in]      config  the lattice configuration the object is currently initialized for
      /// \param[in]      alpha   constant to multiply the coefficient's derivatives by
      /// \param[out]     der_vec vector to which to add the derivative contributions from this object
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void add_derivatives(const pcps::Configuration<S> & config,
                           const pcps::MantExp<S> & alpha,
                           std::vector<S> & der_vec) {

        const S product = ( this->_curr_val * alpha ).get();
        for (int i = _n_fixed; i < _data_pos.size(); i++) {
          assert( _data_pos[i] >= 0 );
          assert( _cmpd_inds[i] >= 0 );
          assert( _data_pos[i] + _cmpd_inds[i] < der_vec.size() );
          der_vec[ _data_pos[i] + _cmpd_inds[i] ] += product;
        }

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief For each orbital pair excitation, contracts the wavefunction's derivatives at the
      ///        corresponding excited configuration with the supplied trial vector.
      ///        It is assumed that the object is initialized for the supplied configuration.
      ///
      /// \param[in]      config      the lattice configuration to excite from.
      /// \param[in]      alpha       constant to multiply the contraction by
      /// \param[in]      trial_vec   the vector to be contracted
      /// \param[in,out]  out_vec_a   vector to which to add the alpha orbital result
      /// \param[in,out]  out_vec_b   vector to which to add the beta orbital result
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void contract_der_to_orb_pairs(const pcps::Configuration<S> & config,
                                     const pcps::MantExp<S> & alpha,
                                     const std::vector<S> & trial_vec,
                                     std::vector<S> & out_vec_a,
                                     std::vector<S> & out_vec_b)
      {
        throw pcps::Exception("NOT YET IMPLEMENTED");
      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief For each wavefunction variable, contracts the various singly excited configurations'
      ///        wavefunction derivatives with respect to that variable with the supplied trial vector.
      ///        It is assumed that the object is initialized for the supplied configuration.
      ///
      /// \param[in]      config        the lattice configuration to excite from.
      /// \param[in]      alpha         constant to multiply the contraction by
      /// \param[in]      trial_vec_a   the vector to be contracted for alpha orbitals
      /// \param[in]      trial_vec_b   the vector to be contracted for beta orbitals
      /// \param[in,out]  out_vec       vector to which to add the result of the contraction
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void contract_orb_pairs_to_der(const pcps::Configuration<S> & config,
                                     const pcps::MantExp<S> & alpha,
                                     const std::vector<S> & trial_vec_a,
                                     const std::vector<S> & trial_vec_b,
                                     std::vector<S> & out_vec)
      {
        throw pcps::Exception("NOT YET IMPLEMENTED");
      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief Multiply the supplied trial vector by the wavefunction derivative outer product matrix
      ///        (the local contribution to the overlap matrix) for the given configuration.
      ///        It is assumed that the object is initialized for the supplied configuration.
      ///
      /// \param[in]      config      the lattice configuration to excite from.
      /// \param[in]      alpha       constant to multiply the contraction by
      /// \param[in]      trial_vec   the vector to be contracted
      /// \param[in,out]  out_vec     vector to which to add the result of the contraction
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void contract_der_to_der(const pcps::Configuration<S> & config,
                               const pcps::MantExp<S> & alpha,
                               const std::vector<S> & trial_vec,
                               std::vector<S> & out_vec)
      {
        throw pcps::Exception("NOT YET IMPLEMENTED");
      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief A class for sorting the correlators within a correlator product.
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      class CorrelatorSorter {

        public:

          typedef typename boost::shared_ptr<pcps::CoeffObj<S> > Ptr;

          bool operator()(const Ptr & lhs, const Ptr & rhs) const {
            assert( typeid(*lhs) == typeid(pcps::Correlator<S>) );
            assert( typeid(*rhs) == typeid(pcps::Correlator<S>) );
            const pcps::Correlator<S> * lhs_corr = dynamic_cast<const pcps::Correlator<S> *>(lhs.get());
            const pcps::Correlator<S> * rhs_corr = dynamic_cast<const pcps::Correlator<S> *>(rhs.get());
            if ( lhs_corr->value_data().fixed() && ! rhs_corr->value_data().fixed() )
              return true;
            return false;
          }

      };

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief Initializes the object for the given lattice configuration.
      ///
      /// \param[in]      config  a lattice configuration
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void initialize(const pcps::Configuration<S> & config) {

        // sort the correlators so that fixed correlators appear first
        std::sort(this->_co_deps->begin(), this->_co_deps->end(), CorrelatorSorter());

        // get the number of fixed correlators
        _n_fixed = 0;
        for (int j = 0; j < this->_co_deps->size(); j++) {
          pcps::Correlator<S> * corr = dynamic_cast<pcps::Correlator<S> *>(this->_co_deps->at(j).get());
          if ( corr->value_data().fixed() )
            _n_fixed++;
        }

        // count how many coefficients touch each lattice site and populate the relavant vectors
        for (int site = 0; site < config.fock_config().size(); site++) {

          // count the number of touching correlators
          int n_touch = 0;
          for (int j = 0; j < this->_co_deps->size(); j++) {
            pcps::Correlator<S> * corr = dynamic_cast<pcps::Correlator<S> *>(this->_co_deps->at(j).get());
            if ( std::set<int>(corr->indices().begin(), corr->indices().end()).count(site) > 0 )
              n_touch++;
          }

          // resize the vectors
          this->_touching_corrs.at(site).resize(n_touch);
          this->_occ_coeffs.at(site).resize(n_touch);

          // record the correlator indices and occupation coefficients
          int k = 0;
          for (int j = 0; j < this->_co_deps->size(); j++) {
            pcps::Correlator<S> * corr = dynamic_cast<pcps::Correlator<S> *>(this->_co_deps->at(j).get());
            if ( std::set<int>(corr->indices().begin(), corr->indices().end()).count(site) > 0 ) {
              this->_touching_corrs.at(site).at(k) = j;
              int coeff = 1;
              std::vector<int>::const_reverse_iterator it = corr->indices().rbegin();
              std::vector<int>::const_reverse_iterator it_end = corr->indices().rend();
              for ( ; it != it_end && *it != site; it++)
                coeff *= 2;
              assert( it != it_end );
              this->_occ_coeffs.at(site).at(k) = coeff;
              k++;
            }
          }
          assert( k == n_touch );

        }

        // initialize the current value and the vectors of compound indeces, data pointers, and data positions
        S sum_total = pcps::zero<S>();
        this->_cmpd_inds.resize(this->_co_deps->size());
        this->_data_ptrs.resize(this->_co_deps->size());
        this->_data_pos.resize(this->_co_deps->size());
        for (size_t i = 0; i < this->_co_deps->size(); i++) {
          pcps::Correlator<S> * corr = dynamic_cast<pcps::Correlator<S> *>(this->_co_deps->at(i).get());
          sum_total += corr->value_data().data_vec().at(0);
          this->_cmpd_inds.at(i) = 0;
          this->_data_ptrs.at(i) = corr->data_ptr();
          this->_data_pos.at(i) = corr->value_data().data_pos();
        }
        this->_curr_val = pcps::MantExp<S>(std::exp(sum_total));

        // initialize the current value for the requested configuration
        for (int site = 0; site < config.fock_config().size(); site++)
          this->update_one_site(site, config.fock_config().at(site), 0, this->_curr_val);

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief Updates the object so that it is initialized for a new configuration that differs from
      ///        the currently initialized configuration as specified.
      ///
      /// \param[in]      new_config  the lattice configuration we want to object to be initialized for
      /// \param[in]      old_config  the lattice configuration the object is currently initialized for
      /// \param[in]      diff        an object specifying how the configurations differ
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void update(const pcps::Configuration<S> & new_config,
                  const pcps::Configuration<S> & old_config,
                  const pcps::ConfigDiff<S> & diff) {

        for (std::vector<int>::const_iterator site = diff.newly_empty().begin(); site != diff.newly_empty().end(); site++)
          this->update_one_site(*site, 0, 1, this->_curr_val);

        for (std::vector<int>::const_iterator site = diff.newly_full().begin(); site != diff.newly_full().end(); site++)
          this->update_one_site(*site, 1, 0, this->_curr_val);

      }

  };

}

#endif
